package com.knxhd.controller;

import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import com.knxhd.enums.Role;
import com.knxhd.pojo.HttpResult;
import com.knxhd.pojo.LlmResponse;
import com.knxhd.pojo.LlmResponse.Choice;
import com.knxhd.pojo.LlmResponse.MessageDetail;
import com.knxhd.pojo.LlmResponse.Usage;

import cn.hutool.core.collection.ListUtil;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.output.TokenUsage;
import jakarta.annotation.Resource;

/**
 * Author: tianluhua tianlh@qdcares.cn
 * Date: 2025-07-12 10:46:57
 * Description: 大模型调用
 */
@RestController
@RequestMapping(value = "/chat")
public class ChatController {

  @Resource(name = "qwen")
  private ChatModel chatModel;

  @Resource(name = "ds")
  private ChatModel dsChatModel;

  /**
   * 模型调用
   * 
   * @param input
   * @return
   */
  @GetMapping(value = "")
  public String chatString(@RequestParam(defaultValue = "你好") String input,
      @RequestParam(defaultValue = "qwen") String type) {
    if ("qwen".equals(type)) {
      return chatModel.chat(input);
    } else if ("ds".equals(type)) {
      return dsChatModel.chat(input);
    }

    return "请选择模型";
  }

  /**
   * 模型调用
   * 
   * @param input
   * @return
   */
  @GetMapping(value = "/completions")
  public HttpResult<LlmResponse> completions(@RequestParam(defaultValue = "你好") String input,
      @RequestParam(defaultValue = "qwen") String type) {
    ChatResponse chatResponse = null;
    if ("qwen".equals(type)) {
      chatResponse = chatModel.chat(UserMessage.from(input));
    } else if ("ds".equals(type)) {
      chatResponse = dsChatModel.chat(UserMessage.from(input));
    }

    if (chatResponse == null) {
      return HttpResult.error(400, "类型不存在");
    }

    AiMessage aiMessage = chatResponse.aiMessage();
    TokenUsage tokenUsage = chatResponse.tokenUsage();
    String text = aiMessage.text();
    return HttpResult.success(LlmResponse.builder()
        .choices(ListUtil.of(new Choice(new MessageDetail(Role.getRole(aiMessage.type().name()), text))))
        .usage(new Usage(tokenUsage.outputTokenCount(), tokenUsage.inputTokenCount(), tokenUsage.totalTokenCount()))
        .build());
  }

  @GetMapping(value = "/chatStream")
  public SseEmitter chatStream(@RequestParam(defaultValue = "你好") String input,
      @RequestParam(defaultValue = "qwen") String type) {

    ChatResponse chatResponse = null;
    if ("qwen".equals(type)) {
      chatResponse = chatModel.chat(UserMessage.from(input));
    } else if ("ds".equals(type)) {
      chatResponse = dsChatModel.chat(UserMessage.from(input));
    }

    SseEmitter emitter = new SseEmitter(5 * 60 * 1000L);
    if (chatResponse == null) {
      return emitter;
    }

    AiMessage aiMessage = chatResponse.aiMessage();
    TokenUsage tokenUsage = chatResponse.tokenUsage();
    String text = aiMessage.text();
    HttpResult.success(LlmResponse.builder()
        .choices(ListUtil.of(new Choice(new MessageDetail(Role.getRole(aiMessage.type().name()), text))))
        .usage(new Usage(tokenUsage.outputTokenCount(), tokenUsage.inputTokenCount(), tokenUsage.totalTokenCount()))
        .build());

    return emitter;
  }
}
